from enum import Enum
from typing import Optional, List
from pydantic import BaseModel, Field
from abc import ABCMeta
from .chat_completion_function import ChatCompletionFunctionCall


__all__ = [
    "ChatCompletionRole",
    "ChatCompletionMessage",
    "ChatCompletionSystemMessage",
    "ChatCompletionUserMessage",
    "ChatCompletionAssistantMessage",
    "ChatCompletionFunctionMessage",
    "ChatCompletionFinishReason",
    "ChatCompletion",
]


class ChatCompletionRole(str, Enum):
    system = "system"
    assistant = "assistant"
    user = "user"
    function = "function"


class ChatCompletionMessage(BaseModel, metaclass=ABCMeta):
    content: Optional[str] = Field(
        None,
        description="The content of the message.",
        examples=["Hello!"],
    )


class ChatCompletionSystemMessage(ChatCompletionMessage):
    role: ChatCompletionRole = Field(
        ...,
        Literal=ChatCompletionRole.system,
        description="The role of the message, which is always 'system' for a system message",
        examples=["system"],
    )


class ChatCompletionUserMessage(ChatCompletionMessage):
    role: ChatCompletionRole = Field(
        ...,
        Literal=ChatCompletionRole.user,
        description="The role of the message, which is always 'user' for a user message",
        examples=["user"],
    )


class ChatCompletionAssistantMessage(ChatCompletionMessage):
    role: ChatCompletionRole = Field(
        ...,
        Literal=ChatCompletionRole.assistant,
        description="The role of the message, which is always 'assistant' for an assistant message",
        examples=["assistant"],
    )

    function_calls: Optional[List[ChatCompletionFunctionCall]] = Field(
        None,
        description="The funcion calls requested by the assistant.",
    )


class ChatCompletionFunctionMessage(ChatCompletionMessage):
    role: ChatCompletionRole = Field(
        ...,
        Literal=ChatCompletionRole.function,
        description="The role of the message, which is always 'function' for a function message",
        examples=["function"],
    )

    id: str = Field(
        ...,
        description="The corresponding id of the tool requested by the assistant.",
        examples=["call_abc123"],
    )


class ChatCompletionFinishReason(str, Enum):
    stop = "stop"
    length = "length"
    function_calls = "function_calls"
    recitation = "recitation"
    error = "error"
    unknown = "unknown"


class ChatCompletion(BaseModel):
    object: str = Field(
        "ChatCompletion",
        Literal="ChatCompletion",
        description="The object type, which is always 'ChatCompletion'.",
        examples=["ChatCompletion"],
    )

    finish_reason: ChatCompletionFinishReason = Field(
        ...,
        description="The reason why the generation is finished.",
        examples=["stop"],
    )

    message: ChatCompletionAssistantMessage = Field(
        ...,
        description="The message generated by the assistant.",
    )

    created_timestamp: int = Field(
        ...,
        description="The timestamp in milliseconds when the response is created.",
        examples=[1700000000000],
    )
